#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import hashlib
import sys
import os


def CalculateMD5(filepath):
    m = hashlib.md5()
    with open(filepath, 'rb') as fp:
        while True:
            blk = fp.read(4096)  # 4KB per block
            if not blk: break
            m.update(blk)
    return m.hexdigest()


def CalculateSHA1(filepath):
    s = hashlib.sha1()
    with open(filepath, 'rb') as fp:
        while True:
            blk = fp.read(4096)  # 4KB per block
            if not blk: break
            s.update(blk)
    return s.hexdigest()


if __name__ == "__main__":
    parent_path = sys.argv[1]

    DB_NAME = "pmc"
    COLLECTION_NAME = "filelist"

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    results = collection.find({"nofile": {"$exists": 0}})
    for result in results:
        id = result["_id"]
        path = result["path"]
        if os.path.exists(parent_path + "/" + path):
            md5 = CalculateMD5(parent_path + "/" + path)
            sha1 = CalculateSHA1(parent_path + "/" + path)
            collection.update_one({"_id": id}, {"$set": {"md5": md5, "sha1": sha1, "nofile": False}})
    connection.close()
